import torch
from torch import nn

class RNN(nn.Module):

    def __init__(self, input_size, hidden_size, output_size):
        """
        构造一个自定义RNN模型，实现短文本分类
        :param input_size: 输入维度
        :param hidden_size: 隐藏层
        :param output_size: 输出维度
        """

        super(RNN, self).__init__()

        self.hidden_size = hidden_size
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        self.linear = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input, hidden):
        output, hidden = self.rnn(input.unsqueeze(0), hidden.unsqueeze(0))
        output = self.linear(output.squeeze(0))
        output = self.softmax(output)
        return output, hidden.squeeze(0)


    def init_hidden(self):
        return torch.zeros(1, self.hidden_size)


if __name__ == '__main__':
    input_size = 768
    hidden_size = 128
    n_categories = 2

    input = torch.rand(1, input_size)
    hidden = torch.rand(1, hidden_size)

    rnn = RNN(input_size, hidden_size, n_categories)
    outputs, hidden = rnn(input, hidden)
    print("outputs:", outputs)
    print("hidden:", hidden)

